import modules.scripts as scripts
import gradio as gr
import os
from modules.paths import models_path
from torch.hub import download_url_to_file
from modules import script_callbacks

scripts_sam_model_dir = os.path.join(scripts.basedir(), "models")
sd_sam_model_dir = os.path.join(models_path, "sam")
sam_model_dir = sd_sam_model_dir if os.path.exists(sd_sam_model_dir) else scripts_sam_model_dir

def get_sam_model_ids():
    """Get SAM model ids list.

    Returns:
        list: SAM model ids list
    """
    sam_model_ids = [
        "sam_vit_h_4b8939.pth",
        "sam_vit_l_0b3195.pth",
        "sam_vit_b_01ec64.pth",
    ]
    return sam_model_ids
def download_model(sam_model_id):
    """Download SAM model.

    Args:
        sam_model_id (str): SAM model id

    Returns:
        str: download status
    """
    url_sam = "https://dl.fbaipublicfiles.com/segment_anything/" + sam_model_id

    sam_checkpoint = os.path.join(sam_model_dir, sam_model_id)
    if not os.path.isfile(sam_checkpoint):
        try:
            download_url_to_file(url_sam, sam_checkpoint)
        except Exception as e:
            print(e)
            return str(e)

        return "Download complete"
    else:
        return "Model already exists"

def on_ui_tabs():
    sam_model_ids = get_sam_model_ids()

    with gr.Blocks(analytics_enabled=False) as ui_component:
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    with gr.Column():
                        sam_model_id = gr.Dropdown(label="Segment Anything Model ID", elem_id="sam_model_id",
                                                   choices=sam_model_ids, show_label=True)
                    with gr.Column():
                        with gr.Row():
                            load_model_btn = gr.Button("Download model", elem_id="load_model_btn")
                        with gr.Row():
                            status_text = gr.Textbox(label="", elem_id="status_text", max_lines=1, show_label=False,
                                                     interactive=False)
            load_model_btn.click(download_model, inputs=[sam_model_id], outputs=[status_text])
        return [(ui_component, "SAM API", "sam_api_tab")]

script_callbacks.on_ui_tabs(on_ui_tabs)
